using System;
using System.Reflection;
using NBody;
using NBody.Analysis;
using NBody.Cosmology;
using OptionParser;
using DataFile;

[assembly: AssemblyTitle ("NBodyDF")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Calculate the distribution function f(E, J).")]
[assembly: AssemblyCopyright ("2006 Joseph D. MacMillan")]

public class DFOptions : InputOutputGetOptions
{
    [Option ("-b", "Number of bins")]
    public int NumBins = 100;
    
    [Option ("-e", "Softening length")]
    public double Softening;
    
    [Option ("-f", "File for potential data")]
    public string PotFile;
}

public class NBodyDF
{
    public static void Main(string[] args)
    {
        // Process command-line options
        DFOptions opts = new DFOptions();
        opts.ProcessArgs(args, "df");
                        
        // Read in a N-body system from the given file
        NBodySystem s = NBodySystem.Read(opts.InFile);
            
        if (opts.PotFile != "")
            Energy.SetPotentialFromData(ref s, Table.Read(opts.PotFile));
        else
            Energy.CalculatePotentialShell(ref s, opts.Softening);
        
        double e_tot;
        Energy.CalculateTotalEnergy(ref s, out e_tot);   
        s.Sort();
        
        // find emin, emax, and jmin, jmax
        double Emin = 1e6, Emax = -1e6, Jmin = 1e10, Jmax = -1e10;
        for (int i = 0; i < s.NumParts; i++)
        {
            EnergyParticle p = (EnergyParticle)s[i];
            double j = p.Mass * Math.Sqrt(p.J2);
            double e = p.Energy;
            if (j > Jmax) Jmax = j;
            if (j < Jmin) Jmin = j;
            if (e > Emax) Emax = e;
            if (e < Emin) Emin = e;
        }
        Console.Error.WriteLine("Binning energy from Emin = {0:f6} to Emax = {1:f6}", Emin, Emax);
        Console.Error.WriteLine("Binning angular momentum from Jmin = {0:f6} to Jmax = {1:f6}", Jmin, Jmax);
        
        double dE = (Emax - Emin) / (double)opts.NumBins;
        double dJ = (Jmax - Jmin) / (double)opts.NumBins;
        
        Table data = new Table(opts.NumBins * opts.NumBins, 5, "Distribution function");
        
        for (int i = 0; i < opts.NumBins; i++)
        {
            double e = Emin + dE * i;
            //FindJMinMax(s, e, e + dE, out Jmin, out Jmax);
            //dJ = (Jmax - Jmin) / (double)opts.NumBins;
            for (int k = 0; k < opts.NumBins; k++)
            {
                double j = Jmin + dJ * k;
                
                // calculate number of particles in this bin
                double N = CalcNumberDensity(s, e, e + dE, j, j + dJ);
                
                // calculate density of states
                double w = CalcDenStates(s, e, j);
        
                double f = N / w;
                if (w == 0.0)
                    f = 0.0;                
                
                data[k + opts.NumBins * i, 0] = e;
                data[k + opts.NumBins * i, 1] = j;
                data[k + opts.NumBins * i, 2] = N;
                data[k + opts.NumBins * i, 3] = w;
                data[k + opts.NumBins * i, 4] = f;
                Console.Error.WriteLine("({0})  J = {1:f6}, N = {2:f6}, w = {3:f6}, f = {4:f6}", k, j, N, w, f);
            }
            Console.Error.WriteLine("({0})  E = {1:f6}, Jmin = {2:f6}, Jmax = {3:f6}", i, e, Jmin, Jmax);
        }
        
        data.Print(opts.OutFile);
        
    }
    
    static void FindJMinMax(NBodySystem s, double emin, double emax, out double jmin, out double jmax)
    {
        jmin = jmax = s[0].Mass * Math.Sqrt(s[0].J2);
        int c = 0;
        for (int i = 1; i < s.NumParts; i++)
        {
            EnergyParticle p = (EnergyParticle)s[i];
            double e = p.Energy;
            double j = p.Mass * Math.Sqrt(p.J2);
            if (e < emax && e > emin)
            {
                c++;
                if (j > jmax) jmax = j;
                if (j < jmin) jmin = j;
            }            
        }
        jmax += 1e-10;
        jmin -= 1e-10;
        Console.Error.WriteLine("**** There are {0} particles between E = {1:f6} and E = {2:f6}", c, emin, emax);
    }
    
    static double CalcNumberDensity(NBodySystem s, double emin, double emax, double jmin, double jmax)
    {
        int n = 0;
        for (int i = 0; i < s.NumParts; i++)
        {
            EnergyParticle p = (EnergyParticle)s[i];
            double e = p.Energy;
            double j = p.Mass * Math.Sqrt(p.J2);
            
            if (e < emax && e > emin && j < jmax && j > jmin)
                n++;
        }
        return (double)n;
    }
    
    static double CalcDenStates(NBodySystem s, double e, double j)
    {
        double[] r = new double[s.NumParts];
        double[] g1 = new double[s.NumParts];
        for (int i = 0; i < s.NumParts; i++)
        {
            r[i] = s[i].Radius;
            g1[i] = 2.0 * (e - ((EnergyParticle)s[i]).Potential) - j * j / r[i] / r[i];
            if (g1[i] < 0.0)
                g1[i] = 0.0;
            else
                g1[i] = 1.0 / Math.Sqrt(g1[i]);
        }
        
        return 8.0 * Math.PI * Math.PI * Integration.Trapezoidal(r, g1, 0, s.NumParts - 1);
    }
}
